#include "SplitFbxByMaterialWithHierarchy.h"
#include <iostream>
#include <algorithm>

FbxNode* fbxTool::SplitFbxByMaterialWithHierarchy::FindNodeByPath(FbxNode* rootNode, const vector<int>& indexPath) {
    FbxNode* currentNode = rootNode;

    // 遍历路径中的每个索引
    for (const auto& index : indexPath) {
        if (index < 0 || index >= currentNode->GetChildCount()) {
            // 如果索引超出范围，返回 nullptr
            return nullptr;
        }

        // 获取当前节点的子节点，使用索引来直接定位子节点
        currentNode = currentNode->GetChild(index);
    }

    // 最终返回匹配的节点
    return currentNode;
}

void fbxTool::SplitFbxByMaterialWithHierarchy::RemoveNode(FbxNode* node, FbxScene* scene) {
    if (node) {
        FbxNode* parent = node->GetParent();
        if (parent)
        {
            if (parent->GetChildCount() == 1) {
                RemoveNode(parent, scene);
            }
            else {
                RemoveDescendants(node, scene);
            }
        }
    }
}

void fbxTool::SplitFbxByMaterialWithHierarchy::RemoveDescendants(FbxNode* node, FbxScene* scene) {
    if (node) {
        // 递归删除该节点的所有子孙节点（不局限于直接子节点）
        int childCount = node->GetChildCount();
        for (int i = 0; i < childCount; ++i) {
            RemoveDescendants(node->GetChild(i), scene);  // 递归删除每个子节点
        }
        FbxNode* parent = node->GetParent();
        scene->RemoveNode(node);
        parent->RemoveChild(node);
        node->Destroy();
        node = nullptr;
    }
}


void fbxTool::SplitFbxByMaterialWithHierarchy::RemoveEmptyNode(FbxNode* node, FbxScene* scene) {
    if (!node) return;

    // 获取节点属性
    FbxNodeAttribute* attr = node->GetNodeAttribute();

    // 如果节点属性为网格类型，加入到meshNodes中
    if (attr && attr->GetAttributeType() == FbxNodeAttribute::eMesh) {
        return;
    }

    // 遍历子节点并递归处理
    size_t numChildren = node->GetChildCount();
    if (numChildren == 0)
        RemoveNode(node, scene);
    for (size_t i = 0; i < numChildren; ++i) {
        // 递归遍历子节点
        RemoveEmptyNode(node->GetChild(static_cast<int>(i)), scene);
    }
}

void fbxTool::SplitFbxByMaterialWithHierarchy::GetNodePath(FbxNode* node, vector<int>& indexPath) {
    if (node == nullptr) return;

    FbxNode* parentNode = node->GetParent();

    // 递归获取父节点的索引路径
    if (parentNode != nullptr) {
        // 递归向上找到根节点
        GetNodePath(parentNode, indexPath);
    }

    // 获取当前节点在父节点的子节点列表中的索引
    if (parentNode != nullptr) {
        int childCount = parentNode->GetChildCount();
        for (int i = 0; i < childCount; ++i) {
            if (parentNode->GetChild(i) == node) {
                // 找到当前节点在父节点的子节点列表中的索引
                indexPath.push_back(i);
                break;
            }
        }
    }
}


vector<string> fbxTool::SplitFbxByMaterialWithHierarchy::Split(const string& inputPath, const string& outputPath, int numParts)
{
    vector<string> result;

    vector<FbxNode*> meshNodes;
    FbxManager* manager; FbxScene* originalScene;
    GetAllMeshNodes(inputPath, meshNodes, manager, originalScene);

    // 按材质分组
    map<string, vector<FbxNode*>> materialGroups;
    for (FbxNode* node : meshNodes) {
        FbxMesh* mesh = (FbxMesh*)node->GetNodeAttribute();
        if (mesh) {
            int materialCount = node->GetMaterialCount();
            for (int i = 0; i < materialCount; ++i) {
                FbxSurfaceMaterial* material = node->GetMaterial(i);
                string materialName = material ? material->GetName() : "UnknownMaterial";
                materialGroups[materialName].push_back(node);
                break;  // 一般情况下，每个节点对应一个材质，若有多个材质可以调整
            }
        }
    }

    // 分割材质节点
    vector<vector<FbxNode*>> partitionedGroups(numParts);

    // 记录已分配的节点
    set<FbxNode*> assignedNodes;

    // 将材质节点分配到分组中
    for (const auto& group : materialGroups) {
        const string& materialName = group.first;
        const vector<FbxNode*>& nodes = group.second;

        for (FbxNode* node : nodes) {
            if (assignedNodes.count(node) == 0) {
                // 寻找最小的分组
                int bestGroup = 0;
                int minGroupSize = static_cast<int>(partitionedGroups[0].size());
                for (int i = 1; i < numParts; ++i) {
                    if (partitionedGroups[i].size() < minGroupSize) {
                        bestGroup = i;
                        minGroupSize = static_cast<int>(partitionedGroups[i].size());
                    }
                }
                partitionedGroups[bestGroup].push_back(node);
                assignedNodes.insert(node);
            }
        }
    }

    // 导出每个分组到独立的 FBX 文件
    for (int i = 0; i < numParts; ++i) {
        FbxManager* managerRepeat; FbxScene* scene;
        ReadFbx(inputPath, managerRepeat, scene);

        // 克隆整个原始场景的所有节点
        FbxNode* rootNode = scene->GetRootNode();

        vector<FbxNode*> waitRemoveNodes;

        // 遍历分组中的节点，找到与其对应的克隆节点
        for (int j = 0; j < partitionedGroups.size(); j++)
        {
            if (j != i)
            {
                for (FbxNode* node : partitionedGroups[j]) {
                    vector<int> indexPath;
                    GetNodePath(node, indexPath);
                    FbxNode* findNode = FindNodeByPath(rootNode, indexPath);
                    waitRemoveNodes.push_back(findNode);
                }
            }
        }
        for (FbxNode* waitRemoveNode : waitRemoveNodes)
        {
            RemoveNode(waitRemoveNode, scene);
        }
        RemoveEmptyNode(rootNode, scene);
        // 导出该分组的FBX文件
        const string filename = ExportChildFbx(partitionedGroups, managerRepeat, scene, outputPath, i);
        if (!filename.empty())
            result.push_back(filename);

        managerRepeat->Destroy();
    }

    manager->Destroy();
    return result;
}

